// 此代码由DeepSeek生成的代码修改而来
#include "vector/iterator.h"

#include <algorithm>
#include <stdexcept>

Iterator::Iterator(int64_t data_size) : gen_(std::random_device{}())
{
    reset(data_size);
}

// 重置索引和游标（全量洗牌）
void Iterator::reset(size_t data_size)
{
    indices_.resize(data_size);
    for (size_t i = 0; i < data_size; ++i)
    {
        indices_[i] = i;
    }
    std::shuffle(indices_.begin(), indices_.end(), gen_);
    cursor_ = 0;
}

// 高频调用优化版本
std::vector<std::vector<double>> Iterator::get_rand_batch(
    const std::vector<std::vector<double>> &data,
    int64_t batch_size)
{
    // 1. 检查输入有效性
    if (data.empty() || batch_size <= 0)
    {
        return {};
    }

    // 2. 检查是否需要重新洗牌（到达数据末尾时）
    if (cursor_ + batch_size > data.size())
    {
        reset(data.size());
    }

    // 3. 填充批次数据
    std::vector<std::vector<double>> batch;
    batch.reserve(batch_size);
    for (; cursor_ < indices_.size() && batch.size() < static_cast<size_t>(batch_size); ++cursor_)
    {
        batch.emplace_back(data[indices_[cursor_]]);
    }

    return batch;
}